#include "libmri.h"

#include <omp.h>
#include <cuda.h>
#include <cufft.h>

void plan_cufft2c (cufftHandle &plan, int xs1, int xs2, int batch)
{
	int n[2] = {xs2, xs1}; 

	cufft (PlanMany (&plan, 2, n, 0,1,0,0,1,0, CUFFT_C2C, batch));
}

void destroy_cufft2c (cufftHandle &plan)
{
	cufft (Destroy (plan));
}

__global__ void cuifft2shift (cplxf *x, int xs1, int xs2)
{
	x += xs1*xs2*blockIdx.x;

	for (int i2 = threadIdx.x; i2 < xs2; i2 += blockDim.x){
		ifftshift(&x [i2*xs1], xs1, 1);
	}

	barrier ();

	for (int i1 = threadIdx.x; i1 < xs1; i1 += blockDim.x){
		ifftshift(&x [i1], xs2, xs1);
	}
}

void cuifft2shift (cplxf *X, int xs1, int xs2, int batch)
{
	cuifft2shift <<<batch, 64>>> (X,xs1,xs2);
}

__global__ void cufft2shift (cplxf *x, int xs1, int xs2)
{
	x += xs1*xs2*blockIdx.x;

	for (int i2 = threadIdx.x; i2 < xs2; i2 += blockDim.x)
		fftshift (&x [i2*xs1], xs1, 1);

	barrier ();

	for (int i1 = threadIdx.x; i1 < xs1; i1 += blockDim.x)
		fftshift (&x [i1], xs2, xs1);

	barrier ();

	for (int i = threadIdx.x; i < xs1*xs2; i += blockDim.x)
		x[i] *= (1.0f / sqrtf(xs1*xs2));
}

void cufft2shift (cplxf *X, int xs1, int xs2, int batch)
{
	cufft2shift <<<batch, 64>>> (X, xs1, xs2);
}

void cufft2c_exec (cufftHandle& plan, int direction,
                   cplxf *X, int xs1, int xs2, int batch)
{
	cuifft2shift (X, xs1, xs2, batch);
	cufft (ExecC2C (plan, (cufftComplex*) X, (cufftComplex*) X, direction));
	cufft2shift (X, xs1, xs2, batch);
}

